import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import pandas as pd
import plotly.io as pio

# 设置全局主题
pio.templates.default = "plotly_white"

# 自定义颜色方案
COLOR_SCHEME = {
    "primary": "#1a3a5f",
    "secondary": "#2c6c9c",
    "accent": "#4da6ff",
    "success": "#4caf50",
    "warning": "#ff9800",
    "danger": "#f44336",
    "male": "#3498db",
    "female": "#e74c3c",
    "survived": "#2ecc71",
    "not_survived": "#e74c3c"
}


def create_survival_chart(data):
    """创建现代化的存活率分析图表"""
    # 按船舱等级分析
    class_survival = data.groupby('Pclass')['Survived'].agg(['mean', 'count']).reset_index()
    class_survival.columns = ['Pclass', 'SurvivalRate', 'Count']
    class_survival['Pclass'] = class_survival['Pclass'].map({1: '一等舱', 2: '二等舱', 3: '三等舱'})

    # 按性别分析
    sex_survival = data.groupby('Sex')['Survived'].agg(['mean', 'count']).reset_index()
    sex_survival.columns = ['Sex', 'SurvivalRate', 'Count']
    sex_survival['Sex'] = sex_survival['Sex'].map({0: '男性', 1: '女性'})

    # 按年龄分组分析
    data['AgeGroup'] = pd.cut(data['Age'], bins=[0, 12, 18, 30, 50, 100],
                              labels=['儿童(0-12)', '青少年(13-18)', '青年(19-30)', '中年(31-50)', '老年(50+)'])
    age_survival = data.groupby('AgeGroup', observed=False)['Survived'].agg(['mean', 'count']).reset_index()
    age_survival.columns = ['AgeGroup', 'SurvivalRate', 'Count']

    # 创建子图 - 改为3列
    fig = make_subplots(
        rows=1, cols=3,
        subplot_titles=('船舱等级存活率', '性别存活率', '年龄组存活率'),
        specs=[[{'type': 'bar'}, {'type': 'pie'}, {'type': 'bar'}]],
        horizontal_spacing=0.1,
        vertical_spacing=0.2
    )

    # 船舱等级柱状图 - 使用渐变色
    fig.add_trace(
        go.Bar(
            x=class_survival['Pclass'],
            y=class_survival['SurvivalRate'],
            text=[f'{rate * 100:.1f}%<br>({count}人)' for rate, count in
                  zip(class_survival['SurvivalRate'], class_survival['Count'])],
            textposition='auto',
            marker_color=[COLOR_SCHEME["primary"], COLOR_SCHEME["secondary"], COLOR_SCHEME["accent"]],
            marker_line_color='rgba(0,0,0,0.3)',
            marker_line_width=1,
            opacity=0.9,
            name='船舱等级'
        ),
        row=1, col=1
    )

    # 性别饼图 - 现代设计
    fig.add_trace(
        go.Pie(
            labels=sex_survival['Sex'],
            values=sex_survival['Count'],
            hole=0.5,
            marker_colors=[COLOR_SCHEME["male"], COLOR_SCHEME["female"]],
            textinfo='percent+label',
            hoverinfo='label+percent+value',
            textposition='inside',
            textfont_size=14,
            pull=[0.02, 0.02],
            name='性别'
        ),
        row=1, col=2
    )

    # 年龄组柱状图 - 使用连续色阶
    fig.add_trace(
        go.Bar(
            x=age_survival['AgeGroup'],
            y=age_survival['SurvivalRate'],
            text=[f'{rate * 100:.1f}%' for rate in age_survival['SurvivalRate']],
            textposition='auto',
            marker_color=px.colors.sequential.Blues_r,
            marker_line_color='rgba(0,0,0,0.3)',
            marker_line_width=1,
            opacity=0.9,
            name='年龄组'
        ),
        row=1, col=3
    )

    # 更新布局
    fig.update_layout(
        title_text='泰坦尼克号乘客存活率多维分析',
        title_font_size=20,
        title_x=0.5,
        height=500,
        showlegend=False,
        margin=dict(l=20, r=20, t=80, b=20),
        plot_bgcolor='rgba(0,0,0,0)',
        paper_bgcolor='rgba(0,0,0,0)',
        font_family="Segoe UI, Microsoft YaHei, sans-serif"
    )

    # 更新坐标轴标签
    fig.update_yaxes(
        title_text="存活率",
        row=1,
        col=1,
        tickformat=".0%",
        gridcolor='rgba(0,0,0,0.1)'
    )
    fig.update_yaxes(
        title_text="存活率",
        row=1,
        col=3,
        tickformat=".0%",
        gridcolor='rgba(0,0,0,0.1)'
    )

    # 添加分隔线
    fig.add_shape(
        type="line",
        x0=0.5, y0=0, x1=0.5, y1=1,
        xref="paper", yref="paper",
        line=dict(color="rgba(0,0,0,0.2)", width=1, dash="dot")
    )
    fig.add_shape(
        type="line",
        x0=-0.02, y0=0.5, x1=1.02, y1=0.5,
        xref="paper", yref="paper",
        line=dict(color="rgba(0,0,0,0.2)", width=1, dash="dot")
    )

    return fig


def create_age_fare_scatter(df):
    """创建现代化的年龄-票价散点图"""
    # 创建散点图 - 使用离散颜色映射
    fig = px.scatter(
        df,
        x="Age",
        y="Fare",
        color=df['Survived'].map({0: '未幸存', 1: '已幸存'}),
        size="FamilySize",
        hover_name="Pclass",
        hover_data=["Sex", "Embarked"],
        title="年龄与票价关系分析",
        labels={'color': '存活状态'},
        color_discrete_map={'未幸存': COLOR_SCHEME["not_survived"], '已幸存': COLOR_SCHEME["survived"]},
        opacity=0.7,
        trendline="lowess",
        trendline_color_override=COLOR_SCHEME["primary"]
    )

    # 更新布局
    fig.update_layout(
        height=500,
        xaxis_title="年龄",
        yaxis_title="票价",
        legend_title_text="存活状态",
        title_font_size=18,
        font_family="Segoe UI, Microsoft YaHei, sans-serif",
        plot_bgcolor='rgba(0,0,0,0)',
        paper_bgcolor='rgba(0,0,0,0)',
        hoverlabel=dict(
            bgcolor="white",
            font_size=12,
            font_family="Segoe UI, Microsoft YaHei, sans-serif"
        )
    )

    # 添加参考线
    fig.add_shape(
        type="line",
        x0=0, y0=df['Fare'].mean(), x1=100, y1=df['Fare'].mean(),
        line=dict(color=COLOR_SCHEME["warning"], width=1, dash="dash"),
        name="平均票价"
    )

    fig.add_shape(
        type="line",
        x0=df['Age'].mean(), y0=0, x1=df['Age'].mean(), y1=600,
        line=dict(color=COLOR_SCHEME["accent"], width=1, dash="dash"),
        name="平均年龄"
    )

    return fig


def create_family_survival(data):
    """创建现代化的家庭规模与存活率关系图"""
    # 添加家庭规模分组
    data['FamilyGroup'] = data['FamilySize'].apply(
        lambda x: '单独出行' if x == 1 else
        '小型家庭(2-4人)' if 2 <= x <= 4 else
        '大型家庭(5人+)'
    )

    # 添加 observed=False 参数
    family_survival = data.groupby('FamilyGroup', observed=False)['Survived'].agg(['mean', 'count']).reset_index()
    family_survival.columns = ['FamilyGroup', 'SurvivalRate', 'Count']

    # 按特定顺序排序
    group_order = ['单独出行', '小型家庭(2-4人)', '大型家庭(5人+)']
    family_survival['FamilyGroup'] = pd.Categorical(
        family_survival['FamilyGroup'],
        categories=group_order,
        ordered=True
    )
    family_survival = family_survival.sort_values('FamilyGroup')

    # 创建柱状图
    fig = go.Figure()

    fig.add_trace(go.Bar(
        x=family_survival['FamilyGroup'],
        y=family_survival['SurvivalRate'],
        text=[f'{rate * 100:.1f}%<br>({count}人)' for rate, count in
              zip(family_survival['SurvivalRate'], family_survival['Count'])],
        textposition='auto',
        marker_color=px.colors.sequential.Blues,
        marker_line_color='rgba(0,0,0,0.3)',
        marker_line_width=1,
        opacity=0.9
    ))

    # 更新布局
    fig.update_layout(
        title='家庭规模对存活率的影响',
        title_font_size=20,
        title_x=0.5,
        height=500,
        xaxis_title="家庭规模",
        yaxis_title="存活率",
        yaxis_tickformat=".0%",
        hovermode="x",
        plot_bgcolor='rgba(0,0,0,0)',
        paper_bgcolor='rgba(0,0,0,0)',
        font_family="Segoe UI, Microsoft YaHei, sans-serif",
        xaxis=dict(
            tickangle=-30
        )
    )

    # 添加网格线
    fig.update_yaxes(
        gridcolor='rgba(0,0,0,0.1)',
        showgrid=True
    )

    return fig